import os
import numpy as np
import shutil

topic = ['IT', '体育', '健康', '军事', '招聘', '教育', '文化', '旅游', '财经']
corpus = []     # 存放data根目录里所有的文件, 而后会在这里面随机选取

choice_docs = []        # 存放最终选取的文件的完整的相对路径
def load_docs(fold_path):
	corpus = [[] for _ in range(0, 9)]
	file_list = os.listdir(fold_path)
	for name in file_list:
		file = os.path.join(fold_path, name)
		if os.path.isfile(file) and (os.path.splitext(name)[1] == '.txt'):
			for i in range(0, len(topic)):
				if name.split('_')[0] == topic[i]:
					corpus[i].append(file)          # 将完整的相对路径写入语料库列表中
	return corpus

def choice(n):
	choice_docs = [[0 for _ in range(0, n)] for _ in range(0, 9)]
	for i in range(0, len(topic)):
		for j in range(0, n):
			index = int(np.random.random() * len(corpus[i]))
			choice_docs[i][j] = corpus[i][index]
			corpus[i].remove(corpus[i][index])            # 删掉, 实现不放回的取样
	# print(choice_docs)
	return choice_docs

# 复制选中的文件到指定目录
def copy_choice_file(target_path):
	# 首先清空指定文件夹
	# shutil.rmtree(target_path)
	for i in range(0, len(choice_docs)):
		for j in range(0, len(choice_docs[i])):
			name = os.path.split(choice_docs[i][j])[1]
			new_file = os.path.join(target_path, name)      # 组装新文件的相对路径
			shutil.copyfile(choice_docs[i][j], new_file)


if __name__ == '__main__':
	corpus = load_docs("../data")       # 加载数据根目录全部文件
	choice_docs = choice(100)            # 为每个主题, 随机选出需要的指定数目n的数据文件
	copy_choice_file('../data/sample')      # 将选择的文件复制进目标目录, 保留原文件名